# -*- coding: utf-8 -*-
"""
document_processor.py (Multi-Page Version with NMS)
---------------------
功能:
1. 遍历指定文件夹中的所有图片文件 (e.g., page_001.jpg, page_002.jpg, ...)。
2. 解析一个包含所有页面答案的 answerlist.txt 文件。
3. 对每一张图片，检测图中的红色方框，并使用非极大值抑制（NMS）合并重叠的框。
4. 从解析好的答案中找到对应页面的答案。
5. 将答案内容进行智能排版，填充到对应的红色方框中，并计算精确物理坐标。
6. 将每页的书写任务分割为上半部分和下半部分。
7. 按顺序逐页调用 robot_writer.py 来执行实际的书写任务，并在每页写完后提示用户换纸。

使用方法:
1. 确保已安装所需库: pip install opencv-python numpy Pillow
2. 确保 robot_writer.py 文件与此脚本在同一目录下。
3. 修改下方 `if __name__ == "__main__"` 中的文件夹和文件路径。
4. 将所有扫描的试卷图片（命名为 page_001.jpg, page_002.jpg ...）放入输入文件夹。
5. 准备好包含所有答案的 answerlist.txt 文件。
6. 运行此脚本: python document_processor.py
"""
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import os
import re
import robot_writer  # 导入机械臂书写程序

# ==================== 配置参数 ====================
# A4纸张尺寸 (mm)，用于像素到毫米的转换
A4_WIDTH_MM = 210.0
A4_HEIGHT_MM = 297.0

# 机械臂单次可书写区域尺寸 (mm)
ROBOT_WRITABLE_HEIGHT_MM = 145.0
ROBOT_WRITABLE_WIDTH_MM = 200.0

# ==================== 新增功能：解析多页答案文件 ====================
# ==================== 功能已修改：解析答案时保留末尾空格和换行 ====================
def parse_answers(answer_list_path):
    """
    解析包含所有页面答案的文本文件。
    文件格式应为: "页码-题号: 答案" 或 "页码-题号： 答案"
    此版本会保留答案末尾的空格，并将文本中的 '\\n' 和 '\\p' 转换为真正的换行符。
    返回一个字典，键为页码(str)，值为该页的答案列表(list of str)。
    """
    if not os.path.exists(answer_list_path):
        raise FileNotFoundError(f"无法找到答案列表文件: {answer_list_path}")

    answers_by_page = {}
    with open(answer_list_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.rstrip('\n').rstrip('\r')
            if not line:
                continue
            
            parts = re.split(r'[:：]', line, maxsplit=1)
            if len(parts) != 2:
                print(f"⚠️ 警告: 无法解析答案行，已跳过 -> '{line}'")
                continue

            key, answer = parts
            key = key.strip()
            answer = answer.lstrip() 

            # 【核心修正】: 将文件中的文本 '\\n' 和 '\\p' 替换为程序可识别的换行符 '\n'
            answer = answer.replace('\\n', '\n').replace('\\p', '\n')

            # 从 "3-1" 中提取页码 "3"
            page_num_match = re.match(r'^\d+', key)
            if not page_num_match:
                print(f"⚠️ 警告: 答案索引 '{key}' 格式不正确，无法提取页码，已跳过。")
                continue
            
            page_num_str = page_num_match.group(0)
            
            if page_num_str not in answers_by_page:
                answers_by_page[page_num_str] = []
            answers_by_page[page_num_str].append(answer)
            
    print("✅ 答案文件解析完成，包含以下页面的答案:", list(answers_by_page.keys()))
    return answers_by_page

# ==================== 功能已重构：支持手动换行(\n)和自动换行 ====================
def wrap_text(text, font, max_width_px, draw):
    """
    将 text 按照 font、max_width_px 分行，返回行列表。
    此版本优先根据文本中的 '\\n' 换行，然后再对过长的单行进行自动换行。
    保留行尾的空格。
    """
    all_lines = []
    # 首先根据用户输入的主动换行符 `\n` 分割文本
    paragraphs = text.split('\n')

    for i, paragraph in enumerate(paragraphs):
        # 对于由 `\n` 分割出的每个段落，应用自动换行逻辑
        
        # 如果段落为空，说明是一个空行（由两个连续的\n产生），直接添加一个空字符串代表换行
        if not paragraph and i > 0:
            all_lines.append("")
            continue

        current_line = ""
        for char in paragraph:
            test_line = current_line + char
            bbox = draw.textbbox((0, 0), test_line, font=font)
            line_width = bbox[2] - bbox[0]

            if line_width <= max_width_px:
                current_line = test_line
            else:
                if current_line:
                    all_lines.append(current_line)
                current_line = char
        
        all_lines.append(current_line)
        
    return all_lines


# ==================== 新增功能：非极大值抑制来合并重叠框 ====================
def non_max_suppression(boxes, iou_threshold=0.6):
    """
    对输入的边界框列表执行非极大值抑制，合并高度重叠的框。
    与传统的NMS不同，这里我们合并框而不是简单删除，以应对同一个框被分割检测的情况。
    
    Args:
        boxes (list): 格式为 [(x, y, w, h), ...] 的框列表。
        iou_threshold (float): IoU阈值，超过此值的框将被合并。

    Returns:
        list: 合并后的框列表。
    """
    if not boxes:
        return []

    # 将 (x, y, w, h) 转换为 (x1, y1, x2, y2)
    rects = np.array([[b[0], b[1], b[0] + b[2], b[1] + b[3]] for b in boxes])
    
    # 计算面积
    areas = (rects[:, 2] - rects[:, 0]) * (rects[:, 3] - rects[:, 1])
    # 按y1坐标排序
    indices = np.argsort(rects[:, 1])

    merged_boxes = []
    while len(indices) > 0:
        last = len(indices) - 1
        i = indices[last]
        
        # 将当前框加入到合并列表中
        current_rect = rects[i]
        indices = np.delete(indices, last)

        # 寻找与当前框高度重叠的其他框
        suppress = [last]
        for pos in range(len(indices)):
            j = indices[pos]
            
            # 计算 IoU
            xx1 = np.maximum(current_rect[0], rects[j][0])
            yy1 = np.maximum(current_rect[1], rects[j][1])
            xx2 = np.minimum(current_rect[2], rects[j][2])
            yy2 = np.minimum(current_rect[3], rects[j][3])

            w = np.maximum(0, xx2 - xx1)
            h = np.maximum(0, yy2 - yy1)
            
            intersection = w * h
            union = areas[i] + areas[j] - intersection
            iou = intersection / union if union > 0 else 0
            
            # 如果 IoU 超过阈值，则合并这两个框
            if iou > iou_threshold:
                # 合并框：取两个框的最大外接矩形
                current_rect[0] = min(current_rect[0], rects[j][0])
                current_rect[1] = min(current_rect[1], rects[j][1])
                current_rect[2] = max(current_rect[2], rects[j][2])
                current_rect[3] = max(current_rect[3], rects[j][3])
                
                # 标记此框，以便后续删除
                suppress.append(pos)
        
        # 从索引中删除已被合并的框
        indices = np.delete(indices, [s for s in suppress if s != last])
        
        # 将合并后的大框 (x1, y1, x2, y2) 转换回 (x, y, w, h)
        merged_w = current_rect[2] - current_rect[0]
        merged_h = current_rect[3] - current_rect[1]
        merged_boxes.append((int(current_rect[0]), int(current_rect[1]), int(merged_w), int(merged_h)))
        
    print(f"📦 非极大值抑制：原始检测到 {len(boxes)} 个框，合并后剩余 {len(merged_boxes)} 个框。")
    return merged_boxes


# ==================== 函数已修改：集成 NMS 逻辑 ====================
def detect_and_plan_writing(image_path,
                            answers_for_page, 
                            output_annotated_path,
                            ui_pic_dir,
                            min_area=500):
    """
    检测红框，排版答案，并生成一个包含所有书写任务的计划列表。
    返回: writing_tasks_list
    """
    cv_img = cv2.imread(image_path)
    if cv_img is None:
        raise FileNotFoundError(f"无法加载图片，请检查路径: {image_path}")
    img_h, img_w = cv_img.shape[:2]

    mm_per_pixel_x = A4_WIDTH_MM / img_w
    mm_per_pixel_y = A4_HEIGHT_MM / img_h
    px_per_mm_y = img_h / A4_HEIGHT_MM

    # Step 1: 检测红色闭合框
    hsv = cv2.cvtColor(cv_img, cv2.COLOR_BGR2HSV)
    lower_red1, upper_red1 = np.array([0, 43, 46]), np.array([10, 255, 255])
    lower_red2, upper_red2 = np.array([156, 43, 46]), np.array([180, 255, 255])
    mask_red = cv2.bitwise_or(cv2.inRange(hsv, lower_red1, upper_red1), cv2.inRange(hsv, lower_red2, upper_red2))
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    mask_clean = cv2.morphologyEx(mask_red, cv2.MORPH_CLOSE, kernel, iterations=2)
    mask_clean = cv2.morphologyEx(mask_clean, cv2.MORPH_OPEN, kernel, iterations=1)
    contours, _ = cv2.findContours(mask_clean, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    raw_boxes = []
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area < min_area: continue
        x, y, w, h = cv2.boundingRect(cnt)
        raw_boxes.append((x, y, w, h))

    # ==================== 新增步骤：应用非极大值抑制合并重叠框 ====================
    boxes = non_max_suppression(raw_boxes, iou_threshold=0.6)
    # ==============================================================================

    if not boxes:
        print("⚠️ 当前页面未检测到红框（或已被NMS过滤），跳过书写任务规划。")
        return []
    
    # 按从上到下，从左到右的顺序对框进行排序
    boxes_sorted = sorted(boxes, key=lambda b: (b[1], b[0]))

    # ==================== 代码: 保存含绿色框和蓝色坐标的图片 (1.png) ====================
    img_with_boxes_only = cv_img.copy()
    
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.5
    font_color = (255, 0, 0) # Blue
    thickness = 1
    box_color = (0, 255, 0) # Green
    box_thickness = 2
    
    for (x_px, y_px, w_px, h_px) in boxes_sorted:
        cv2.rectangle(img_with_boxes_only, (x_px, y_px), (x_px + w_px, y_px + h_px), box_color, box_thickness)
        tl, tr, bl, br = (x_px, y_px), (x_px + w_px, y_px), (x_px, y_px + h_px), (x_px + w_px, y_px + h_px)
        tl_text, tr_text, bl_text, br_text = f"({tl[0]},{tl[1]})", f"({tr[0]},{tr[1]})", f"({bl[0]},{bl[1]})", f"({br[0]},{br[1]})"
        (tl_w, tl_h), _ = cv2.getTextSize(tl_text, font, font_scale, thickness)
        (tr_w, tr_h), _ = cv2.getTextSize(tr_text, font, font_scale, thickness)
        (bl_w, bl_h), _ = cv2.getTextSize(bl_text, font, font_scale, thickness)
        (br_w, br_h), _ = cv2.getTextSize(br_text, font, font_scale, thickness)
        cv2.putText(img_with_boxes_only, tl_text, (tl[0], tl[1] - 5), font, font_scale, font_color, thickness)
        cv2.putText(img_with_boxes_only, tr_text, (tr[0] - tr_w, tr[1] - 5), font, font_scale, font_color, thickness)
    
        cv2.putText(img_with_boxes_only, bl_text, (bl[0], bl[1] + bl_h + 5), font, font_scale, font_color, thickness)
        cv2.putText(img_with_boxes_only, br_text, (br[0] - br_w, br[1] + br_h + 5), font, font_scale, font_color, thickness)

    path_1_png = os.path.join(ui_pic_dir, "1.png")
    cv2.imwrite(path_1_png, img_with_boxes_only)
    print(f"✅ (UI) 已将带坐标的红框识别图覆盖保存至: {path_1_png}")
    # =======================================================================================

    if len(answers_for_page) < len(boxes_sorted):
        print(f"⚠️ 警告: 答案数量 ({len(answers_for_page)}) 少于检测到的框数量 ({len(boxes_sorted)})。将只填充部分框。")
    
    answers = answers_for_page[:len(boxes_sorted)]

    annotated_cv = cv_img.copy()
    pil_img = Image.fromarray(cv2.cvtColor(annotated_cv, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(pil_img)
    font_path = r"C:\Windows\Fonts\simfang.ttf"
    if not os.path.exists(font_path):
        raise FileNotFoundError(f"仿宋体字体文件不存在: {font_path}")
    
    all_writing_tasks = []
    print("正在为当前页面规划书写任务...")
    for idx, (x_px, y_px, w_px, h_px) in enumerate(boxes_sorted, start=1):
        if idx > len(answers): break
        answer_text = answers[idx - 1]
        
        w_mm, h_mm = w_px * mm_per_pixel_x, h_px * mm_per_pixel_y
        max_font_mm = min(10.0, h_mm, w_mm / len(answer_text) * 1.5 if answer_text else 10.0)
        min_font_mm = 5.0
        if h_mm < min_font_mm: min_font_mm = h_mm
        max_font_px = max(int(max_font_mm * px_per_mm_y), 1)
        min_font_px = max(int(min_font_mm * px_per_mm_y), 1)
        text_box_w, text_box_h = w_px - 4, h_px - 4
        font_px, final_lines, final_font, line_height_px = max_font_px, [], None, 0
        while font_px >= min_font_px:
            font = ImageFont.truetype(font_path, font_px)
            lines_wrap = wrap_text(answer_text, font, text_box_w, draw)
            bbox = draw.textbbox((0, 0), "汉", font=font)
            lh = bbox[3] - bbox[1]
            if lh * len(lines_wrap) <= text_box_h:
                final_lines, final_font, line_height_px = lines_wrap, font, lh
                break
            font_px -= 1
        if final_font is None:
            font_px = min_font_px
            final_font = ImageFont.truetype(font_path, font_px)
            final_lines = wrap_text(answer_text, final_font, text_box_w, draw)
            bbox = draw.textbbox((0, 0), "汉", font=final_font)
            line_height_px = bbox[3] - bbox[1]
        
        max_line_w_px = max(draw.textbbox((0,0), line, font=final_font)[2] for line in final_lines) if final_lines else 0
        total_text_h_px = line_height_px * len(final_lines)
        x_start_px = x_px + (w_px - max_line_w_px) / 2
        y_start_px = y_px + (h_px - total_text_h_px) / 2
        char_height_mm = (final_font.size / px_per_mm_y) * 0.75

        for i, line in enumerate(final_lines):
            y_line_px = y_start_px + i * line_height_px
            task = {
                "text": line,
                "a4_x_mm": x_start_px * mm_per_pixel_x,
                "a4_y_mm": y_line_px * mm_per_pixel_y,
                "char_height_mm": char_height_mm,
                "char_spacing_ratio": 1.2
            }
            all_writing_tasks.append(task)
            draw.text((x_start_px, y_line_px), line, font=final_font, fill=(0, 0, 0))

        # Re-draw the box on the final annotated image
        cv2.rectangle(annotated_cv, (x_px, y_px), (x_px + w_px, y_px + h_px), (0, 255, 0), 2)
        cv2.putText(annotated_cv, f"{idx}", (x_px, y_px - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,0,0), 2)

    annotated_bgr = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
    
    print(f"✅ 规划完成，共生成 {len(all_writing_tasks)} 行书写任务。")
    print(f"✅ 预览图已生成，将保存至: {output_annotated_path}")
    cv2.imwrite(output_annotated_path, annotated_bgr)
    
    path_2_png = os.path.join(ui_pic_dir, "2.png")
    cv2.imwrite(path_2_png, annotated_bgr)
    print(f"✅ (UI) 已将排版结果图覆盖保存至: {path_2_png}")

    return all_writing_tasks


# ==================== 主程序已重构为多页处理流程 ====================
if __name__ == "__main__":
    # ==================== 请修改以下路径 ====================
    scanned_pages_dir = r"E:\沙粒云\自媒体\2025视频制作\20250610ai高考\流程\anwser\data\exam"
    answer_list_path = r"E:\沙粒云\自媒体\2025视频制作\20250610ai高考\流程\anwser\data\anwser\mapping_result_robust_v2.txt"
    output_dir = r"E:\沙粒云\自媒体\2025视频制作\20250610ai高考\流程\anwser\anotation\annotated_output"
    ui_pic_dir = r"E:\沙粒云\自媒体\2025视频制作\20250610ai高考\流程\aigaokao_ui\pic"
    # =======================================================
    
    if not os.path.exists(output_dir): os.makedirs(output_dir)
    if not os.path.exists(ui_pic_dir): os.makedirs(ui_pic_dir)

    all_answers = parse_answers(answer_list_path)

    try:
        jpg_files = [f for f in os.listdir(scanned_pages_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
        image_files = sorted(jpg_files, key=lambda f: int(re.search(r'(\d+)', f).group(1)))
    except FileNotFoundError:
        print(f"❌ 错误: 找不到输入的图片文件夹: {scanned_pages_dir}")
        exit()
    except (AttributeError, ValueError):
        print(f"❌ 错误: 文件夹中的某些 .jpg 文件名不包含数字，无法进行数字排序。请检查文件命名。")
        exit()

    if not image_files:
        print(f"❌ 错误: 输入文件夹 {scanned_pages_dir} 中没有找到任何图片文件。")
        exit()

    for i, filename in enumerate(image_files):
        print("\n" + "="*20 + f" 开始处理第 {i+1}/{len(image_files)} 页: {filename} " + "="*20)
        
        page_num_match = re.search(r'\d+', filename)
        if not page_num_match:
            print(f"⏩ 文件名 {filename} 不包含数字页码，跳过。")
            continue
        page_num_str = str(int(page_num_match.group(0)))

        answers_for_this_page = all_answers.get(page_num_str, [])
        if not answers_for_this_page:
            print(f"⏩ 页面 {page_num_str} ({filename}) 在答案文件中没有对应答案，跳过。")
            continue

        current_image_path = os.path.join(scanned_pages_dir, filename)
        output_path = os.path.join(output_dir, f"annotated_{filename}")

        writing_tasks = detect_and_plan_writing(current_image_path, answers_for_this_page, output_path, ui_pic_dir)

        if writing_tasks:
            part1_tasks = []
            part2_tasks = []
            for task in writing_tasks:
                if task['a4_x_mm'] < ROBOT_WRITABLE_WIDTH_MM:
                    if task['a4_y_mm'] < ROBOT_WRITABLE_HEIGHT_MM:
                        part1_tasks.append(task)
                    elif ROBOT_WRITABLE_HEIGHT_MM <= task['a4_y_mm'] < 2 * ROBOT_WRITABLE_HEIGHT_MM:
                        task['a4_y_mm'] -= ROBOT_WRITABLE_HEIGHT_MM
                        part2_tasks.append(task)
            
            print("-" * 30)
            print(f"页面 {page_num_str} 书写任务分配：")
            print(f"第 1 部分 (上半页): {len(part1_tasks)} 行")
            print(f"第 2 部分 (下半页): {len(part2_tasks)} 行")
            print("-" * 30)

            try:
                # This assumes robot_writer.py has this function defined
                robot_writer.execute_full_writing_process(part1_tasks, part2_tasks)
            except NameError:
                 print("⚠️ 警告: 'robot_writer' 模块或 'execute_full_writing_process' 函数未定义。跳过实际书写。")
            except Exception as e:
                print(f"❌ 在执行第 {page_num_str} 页书写过程中发生严重错误: {e}")
        else:
            print(f"⏩ 页面 {page_num_str} ({filename}) 未生成任何书写任务。")

        if i < len(image_files) - 1:
            input(f"\n✅✅✅ 页面 {page_num_str} 处理完成。请更换为下一张试卷，然后按 Enter 键继续...")
        else:
            print("\n🎉🎉🎉 所有页面均已处理完毕！")